from sqlalchemy import insert, select

from owl_admin.ext import db
from owl_ai.domain.chat_entity import ChatMessage, Chat
from owl_ai.domain.po import AIChatMessagePO, AIChatPO
from owl_common.sqlalchemy.model import ColumnEntityList


class ChatMapper:
    default_fields = ['id', 'name', 'agent_id', 'create_time', 'update_time', 'create_by', 'update_by']
    default_columns = ColumnEntityList(AIChatPO, default_fields, False)

    @classmethod
    def insert_chat(cls, chat: ChatMessage):
        fields = ['id', 'name', 'agent_id', 'create_time', 'update_time','create_by', 'update_by']
        data = chat.model_dump(include=fields, exclude_none=True, exclude_unset=True)
        stmt = insert(AIChatPO).values(data)
        out = db.session.execute(stmt).inserted_primary_key
        return out[0] if out else 0

    @classmethod
    def find_chat_by_id(cls, chat_id: int):
        stmt = select(*cls.default_columns).where(AIChatPO.id == chat_id)
        row = db.session.execute(stmt).first()
        return cls.default_columns.cast(row, Chat) if row else None

    @classmethod
    def find_chat_list(cls, chat: Chat):
        criteria = []
        if chat.create_by:
            criteria.append(AIChatPO.create_by == chat.create_by)
        stmt = select(*cls.default_columns).where(*criteria).order_by(AIChatPO.update_time.desc())
        rows = db.session.execute(stmt).all()
        return [cls.default_columns.cast(row, Chat) for row in rows] if rows else None


class ChatMessageMapper:
    default_fields = ['id', 'chat_id', 'role', 'content', 'files', 'create_time', 'update_time',
                      'create_user', 'update_user']

    default_columns = ColumnEntityList(AIChatMessagePO, default_fields, False)

    @classmethod
    def insert_message(cls, chat_message: ChatMessage):
        fields = ['id', 'chat_id', 'role', 'content', 'files', 'create_time', 'update_time',
                  'create_user', 'update_user']
        data = chat_message.model_dump(include=fields, exclude_none=True, exclude_unset=True)
        stmt = insert(AIChatMessagePO).values(data)
        out = db.session.execute(stmt).inserted_primary_key
        return out[0] if out else 0

    @classmethod
    def find_messages_by_chat_id(cls, chat_id: int):
        criteria = [AIChatMessagePO.chat_id == chat_id]

        stmt = select(*cls.default_columns).where(*criteria)
        rows = db.session.execute(stmt).all()
        return [cls.default_columns.cast(row, ChatMessage) for row in rows] if rows else None
